import shutil

from util.vgrid_util import getDiffBrightnessSequence, VGridParam, IQMsg
from util.dir_util import makesureDir
from util.HDR import exposure_fusion
from PyVGrid.vgrid import VGrid, BodyPart
from glob import glob
import os.path as osp
from tqdm import tqdm
import numpy as np

if __name__ == '__main__':
    base_path = r"C:\Users\15519\Desktop\不同部位测试图像\胸\org"
    save_path = r"C:\Users\15519\Desktop\不同部位测试图像\胸"

    makesureDir(save_path)                          # 确保存储路径
    # 获得所有的raw文件  一级目录
    name_list = [path.split(".raw")[0] for path in glob(osp.join(base_path, "**", "*.raw"), recursive=True)]
    name_list.sort()
    # 申请tqdm_bar vgridHandle iqHandle
    tqdm_bar = tqdm(enumerate(name_list), total=len(name_list))
    vgridHandle = VGrid()
    iqHandle = IQMsg()
    body_part = BodyPart()
    # 创建特异性的vgrid参数
    vgParam = VGridParam()
    vgParam.PContrast, vgParam.DContrast, vgParam.BContrast, vgParam.brightness = 3, 3.5, 3.5, 0.6
    vgParam.SNRFlag, vgParam.lpSubFlag, vgParam.histFlag = True, True, False
    vgParam.acqMode, vgParam.acqChildmode, vgParam.bodyPart = 1, 0, 5

    for ii, name in tqdm_bar:
        iq_path = name + ".iq"
        raw_path = name + ".raw"
        base_name = osp.basename(name)
        if not osp.exists(raw_path) or not osp.exists(iq_path):
            print(f"have not paired raw or iq file, {name}")
            continue
        iqHandle.parse(iq_path)
        height, width, frame_count = iqHandle.height, iqHandle.width, iqHandle.Count              # 获取尺寸
        vgridHandle.Initialization(width, height, 16)           # 初始化vgrid
        tqdm_bar.set_description(base_name)
        imgs = np.fromfile(raw_path, dtype=np.uint16).reshape((-1, height, width))
        diff_bright_list = getDiffBrightnessSequence(imgs, [0.3, 0.4, 0.54, 0.6, 0.7], vgridHandle, tqdm_bar, vgParam)

        fused_imgs = []
        for frame_id in range(frame_count):
            b_sequence = diff_bright_list[:, frame_id, :, :]
            fused_results = exposure_fusion(b_sequence, refer_illumination_id=2, sigma=0.1, layers_num=7, boday_part=body_part.thorax)
            fused_imgs.append(fused_results)
            tqdm_bar.set_postfix_str(f"fused img {frame_id+1}/{frame_count}")
        fused_imgs = np.array(fused_imgs, dtype=np.uint16)
        fused_imgs.tofile(osp.join(save_path, f"{base_name}.raw"))
        shutil.copy(iq_path, save_path)




